import cv2
import numpy as np

# 读取两帧图像
prev_img = cv2.imread('./modi_tracker/shoulders_frames/frame_0001.png', cv2.IMREAD_GRAYSCALE)
next_img = cv2.imread('./modi_tracker/shoulders_frames/frame_0003.png', cv2.IMREAD_GRAYSCALE)

# 使用Lucas-Kanade方法计算光流
lk_params = dict(winSize=(15, 15), maxLevel=2,
                 criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))

# 计算特征点的光流
feature_params = dict(maxCorners=100, qualityLevel=0.3, minDistance=7, blockSize=7)
p0 = cv2.goodFeaturesToTrack(prev_img, mask=None, **feature_params)
# 将特征点转换为整数类型并调整形状
    
# 计算光流
p1, st, err = cv2.calcOpticalFlowPyrLK(prev_img, next_img, p0, None, **lk_params)

# 绘制光流效果
for i, (new, old) in enumerate(zip(p1, p0)):
    c, d = old.ravel().astype(int)
    a, b = new.ravel().astype(int)
    cv2.arrowedLine(next_img, (c, d), (a, b), (255, 0, 0), 1, tipLength=0.5)

# 显示图像
cv2.imshow('Optical Flow', next_img)
cv2.waitKey(0)
cv2.destroyAllWindows()